from glob import glob
from xml_tool import *
from PIL import Image, ImageDraw, ImageFont
import mmcv
import os, shutil

'''
CASSCDB数据集的读取示例
Created by jiangyizhang
'''

FILE_PATH = "dataset_CASSCDB/"


def use_CASSCDB():
    _init_var()
    seal_xmls = glob(SEAL_YINYANG_XML_FILE + '*')
    for index, seal_xml_path in enumerate(seal_xmls):
        print("处理进度 [%d / %d]" % (index, len(seal_xmls)))
        # 印章逻辑
        seal_image_bean = xml_to_image_bean(seal_xml_path)
        seal_bboxs = get_bboxs_from_xml(seal_xml_path)
        seal_yinyang_list = [object.name for object in seal_image_bean.objects]
        seal_bboxs_yin = []
        seal_bboxs_yang = []
        for seal_bbox_i, seal_bbox in enumerate(seal_bboxs):
            yinyang_label = seal_yinyang_list[seal_bbox_i]
            if yinyang_label == 'yin':
                seal_bboxs_yin.append(seal_bbox)
            else:
                seal_bboxs_yang.append(seal_bbox)
        seal_bboxs_yin = np.asarray(seal_bboxs_yin)
        seal_bboxs_yang = np.asarray(seal_bboxs_yang)

        # 字符逻辑
        char_xml_path = seal_xml_path.replace(SEAL_YINYANG_XML_FILE[:SEAL_YINYANG_XML_FILE.rfind('/')],
                                              CHAR_XML_FILE[:CHAR_XML_FILE.rfind('/')])
        char_image_bean = xml_to_image_bean(char_xml_path)
        char_bboxs = get_bboxs_from_xml(char_xml_path)
        char_name_list = [object.name for object in char_image_bean.objects]

        # 图像
        img_filename_from_xml = seal_image_bean.filename
        img_real_path = IMG_FILE + img_filename_from_xml
        img = Image.open(img_real_path)
        img_np = np.asarray(img)
        img_np_with_bbox = img_np

        # 画框
        if len(seal_bboxs) != 0:
            # 画印章框
            if len(seal_bboxs_yin) != 0:
                # 阴
                img_np_with_bbox = mmcv.imshow_det_bboxes(
                    img_np_with_bbox,
                    np.asarray(seal_bboxs_yin),
                    labels=np.asarray([i for i, _ in enumerate(seal_bboxs_yin)]),
                    # bbox_color="blue",
                    bbox_color="magenta",
                    thickness=8,
                    class_names=["" for i, _ in enumerate(seal_bboxs_yin)],
                    font_scale=3,
                    show=False)
            if len(seal_bboxs_yang) != 0:
                # 阳
                img_np_with_bbox = mmcv.imshow_det_bboxes(
                    img_np_with_bbox,
                    np.asarray(seal_bboxs_yang),
                    labels=np.asarray([i for i, _ in enumerate(seal_bboxs_yang)]),
                    # bbox_color="blue",
                    bbox_color="blue",
                    thickness=8,
                    class_names=["" for i, _ in enumerate(seal_bboxs_yang)],
                    font_scale=3,
                    show=False)

        if len(char_bboxs) != 0:
            # 画字符框
            img_np_with_bbox = mmcv.imshow_det_bboxes(
                img_np_with_bbox,
                char_bboxs,
                labels=np.asarray([i for i, _ in enumerate(char_bboxs)]),
                class_names=["" for i, _ in enumerate(char_bboxs)],
                thickness=3,
                font_scale=3,
                show=False)
        img_with_bbox = Image.fromarray(img_np_with_bbox)

        # 往图像上写中文字符
        draw = ImageDraw.Draw(img_with_bbox)
        ttf = "heiti.ttf"
        ttfront = ImageFont.truetype(ttf, 50)  # 字体大小
        for char_i, char_bbox in enumerate(char_bboxs):
            draw.text((char_bbox[0], char_bbox[1]), char_name_list[char_i], fill=(0, 255, 0), font=ttfront)
        img_filename_with_bboxs = IMG_WITH_BBOXS_FILE + img_filename_from_xml
        img_with_bbox.save(img_filename_with_bboxs)


def _init_var():
    global IMG_FILE, IMG_WITH_BBOXS_FILE, CHAR_XML_FILE, SEAL_YINYANG_XML_FILE
    IMG_FILE = FILE_PATH + "img/"
    IMG_WITH_BBOXS_FILE = FILE_PATH + "img_with_bboxs/"
    CHAR_XML_FILE = FILE_PATH + "char_xml/"
    SEAL_YINYANG_XML_FILE = FILE_PATH + "seal_yinyang_xml/"
    if os.path.exists(os.path.join(IMG_WITH_BBOXS_FILE)):
        shutil.rmtree(IMG_WITH_BBOXS_FILE)
    os.makedirs(IMG_WITH_BBOXS_FILE)


if __name__ == '__main__':
    use_CASSCDB()
